import asyncio
import json
import websockets
from aioredis import Channel, create_redis_pool


redis_conn = {}


async def subscribe_to_redis():
    # conn = await create_redis_pool('redis://172.16.0.130:6379/3', password="abc!@#2021")
    # conn = await create_redis_pool('redis://127.0.0.1:6379/3', password="abc!@#2021")
    conn = await create_redis_pool('redis://172.16.0.71:6379/3', password="abc!@#2021")
    ch1, ch2 = await conn.subscribe('channel:1', 'channel:2')
    assert isinstance(ch1, Channel)
    global redis_conn
    redis_conn = {"channel": ch1, "conn": conn}


websocket_conn = {}


async def browser_server(websocket, path):
    print("-------------", path, websocket)
    channel = redis_conn["channel"]
    conn = redis_conn['conn']
    global websocket_conn
    websocket_conn[path.replace("/", "")] = websocket
    print(websocket_conn)

    # 启动心跳定时器
    heartbeat_task = asyncio.ensure_future(send_heartbeat(websocket))

    try:
        while True:
            # Wait until data is published to this channel
            message = await channel.get()
            # message: df759389cc1e4061b7c7a7aef04eb281:1234
            message = message.decode("utf-8")
            json_message = json.loads(message)
            user_only = json_message.pop("user_only", "")
            # Send unicode decoded data over to the websocket client
            if websocket_conn.get(user_only):
                print(json_message)
                await websocket_conn[user_only].send(json.dumps(json_message))

            # 重置心跳定时器
            heartbeat_task.cancel()
            heartbeat_task = asyncio.ensure_future(send_heartbeat(websocket))
    except websockets.exceptions.ConnectionClosed:
        # 取消心跳定时器
        heartbeat_task.cancel()
        # Free up channel if websocket goes down
        # await conn.execute_pubsub('unsubscribe', channel)
        websocket_conn.pop(path.replace("/", ""), "")
        print("断开连接： ", path.replace("/", ""))


# 发送心跳消息
async def send_heartbeat(websocket):
    while True:
        await websocket.send("heartbeat")
        await asyncio.sleep(10)  # 设置心跳时间间隔（秒）


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    loop.set_debug(True)
    ws_server = websockets.serve(browser_server, '0.0.0.0', 8767)
    loop.run_until_complete(subscribe_to_redis())
    loop.run_until_complete(ws_server)
    loop.run_forever()
